import time
import cv2
import insightface
import numpy as np

model = insightface.app.FaceAnalysis(name='buffalo_l')
model.prepare(ctx_id=0)

cap = cv2.VideoCapture(0)
time.sleep(2)

embeddings = []
max_samples = 10
collected = 0

print("请保持正脸，采集人脸特征中...")

while True:
    ret, frame = cap.read()
    if not ret:
        print("❌ 无法读取摄像头")
        break

    faces = model.get(frame)

    if len(faces) > 0:
        face = faces[0]
        bbox = face.bbox.astype(int)
        cv2.rectangle(frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), 2)

        embedding = face.embedding
        embeddings.append(embedding)
        collected += 1

        cv2.putText(frame, f"采集中 {collected}/{max_samples}", (bbox[0], bbox[1] - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
    else:
        cv2.putText(frame, "未检测到人脸", (20, 40),
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)

    cv2.imshow("Camera", frame)

    if collected >= max_samples:
        break

    if cv2.waitKey(1) & 0xFF == ord('q'):
        print("退出采集")
        break

if embeddings:
    # 计算平均embedding
    my_embedding = np.mean(np.array(embeddings), axis=0)
    np.save("my_face.npy", my_embedding)
    print(f"✅ 采集完成，保存了 {len(embeddings)} 张人脸特征，已保存到 my_face.npy")
else:
    print("❌ 没有采集到有效人脸特征")

cap.release()
cv2.destroyAllWindows()
